import asyncio
import importlib
from typing import Any, Dict
import traceback

from app.core.constants import ExecutionStatus
from app.core.task.executor.base_executor import BaseExecutor
from app.core.task.task_log import scheduler_logger
from app.core.task.store.redis import RedisStore
from app.schemas.job import SchedulerJob, RedisJobExecutionResult
from app.utils.timezone import now


class ModuleExecutor(BaseExecutor):
    """模块任务执行器"""

    def __init__(self, store: RedisStore):
        super().__init__(store)
        self._running_tasks: Dict[str, asyncio.Task] = {}  # 存储正在运行的任务
        self._task_results: Dict[str, RedisJobExecutionResult] = {}  # 存储任务执行结果

    def _create_execution_result(self, job_id: str) -> RedisJobExecutionResult:
        """创建执行结果对象"""
        start_time = now()
        return RedisJobExecutionResult(
            job_id=job_id, start_time=start_time, end_time=start_time, duration=0.0  # 初始化时设置相同的时间
        )

    def _update_execution_result(
        self, result: RedisJobExecutionResult, status: ExecutionStatus, execution_result: Any = None, error: Exception = None
    ) -> RedisJobExecutionResult:
        """更新执行结果"""
        result = super()._update_execution_result(result, status, execution_result, error)
        if error is not None and hasattr(error, "__traceback__"):
            result.traceback = "".join(traceback.format_tb(error.__traceback__))
        return result

    async def cancel_job(self, job_id: str) -> bool:
        """
        取消正在执行的任务

        Args:
            job_id: 任务ID

        Returns:
            bool: 是否成功取消
        """
        if job_id in self._running_tasks:
            task = self._running_tasks[job_id]
            if not task.done():
                # 创建取消结果
                execution_result = self._create_execution_result(job_id)
                execution_result = self._update_execution_result(
                    execution_result, ExecutionStatus.CANCELLED, error=Exception("任务被手动取消")
                )
                # 存储结果
                await self.store.add_job_execution_log(job_id, execution_result)
                self._task_results[job_id] = execution_result

                # 取消任务
                task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    scheduler_logger.info(f"任务 {job_id} 已被取消")
                    return True
                except Exception:
                    pass
            del self._running_tasks[job_id]
        return False

    async def execute_job(self, job: SchedulerJob) -> RedisJobExecutionResult:
        """执行任务并存储结果"""
        job_id = job.id
        execution_result = self._create_execution_result(job_id)
        task = None

        try:
            module = importlib.import_module(job.module_path)
            func = getattr(module, job.function_name)

            task = asyncio.create_task(func(*job.args, **job.kwargs))
            # 等待任务真正开始执行
            await asyncio.sleep(0.001)  # 给予一点时间让任务开始执行
            self._running_tasks[job_id] = task  # 记录运行中的任务

            try:
                result = await asyncio.wait_for(task, timeout=job.timeout)
                execution_result = self._update_execution_result(
                    execution_result, ExecutionStatus.SUCCESS, execution_result=result
                )
                scheduler_logger.info(f"任务 {job_id} 执行成功, 耗时: {execution_result.duration:.2f}秒")
            except asyncio.TimeoutError as e:
                execution_result = self._update_execution_result(
                    execution_result,
                    ExecutionStatus.TIMEOUT,
                    error=Exception(f"TimeoutError: 任务执行超时 (>{job.timeout}秒)"),
                )
                scheduler_logger.error(
                    f"任务 {job_id} 执行超时 (>{job.timeout}秒), 耗时: {execution_result.duration:.2f}秒"
                )
                if task and not task.done():
                    task.cancel()
                    try:
                        await task
                    except asyncio.CancelledError:
                        pass
            except asyncio.CancelledError:
                # 如果任务被取消，返回已存储的取消结果
                if job_id in self._task_results:
                    return self._task_results[job_id]
                # 如果没有存储的结果，创建一个新的取消结果
                execution_result = self._update_execution_result(
                    execution_result, ExecutionStatus.CANCELLED, error=Exception("任务被手动取消")
                )
                scheduler_logger.info(f"任务 {job_id} 被取消, 耗时: {execution_result.duration:.2f}秒")
                raise  # 重新抛出取消异常

        except asyncio.CancelledError:
            # 如果任务被取消，返回已存储的取消结果
            if job_id in self._task_results:
                return self._task_results[job_id]
            # 如果没有存储的结果，创建一个新的取消结果
            execution_result = self._update_execution_result(
                execution_result, ExecutionStatus.CANCELLED, error=Exception("任务被手动取消")
            )
            scheduler_logger.info(f"任务 {job_id} 被取消, 耗时: {execution_result.duration:.2f}秒")
            raise  # 重新抛出取消异常

        except Exception as e:
            execution_result = self._update_execution_result(execution_result, ExecutionStatus.FAILED, error=e)
            scheduler_logger.error(
                f"任务 {job_id} 执行失败, "
                f"耗时: {execution_result.duration:.2f}秒\n"
                f"错误: {execution_result.error_message}\n"
                f"堆栈: {execution_result.traceback}"
            )

        finally:
            if job_id in self._running_tasks:
                del self._running_tasks[job_id]
            # 不要在这里删除 _task_results，让它在下一次任务执行时自然清理
            await self.store.add_job_execution_log(job_id, execution_result)
            return execution_result
